package cn.org.wenet.wekws;

import androidx.appcompat.app.AppCompatActivity;
import androidx.core.app.ActivityCompat;
import androidx.core.content.ContextCompat;

import android.Manifest;
import android.content.Context;
import android.content.pm.PackageManager;
import android.content.res.AssetManager;
import android.media.AudioFormat;
import android.media.AudioRecord;
import android.media.MediaRecorder;
import android.os.Bundle;
import android.os.Process;
import android.util.Log;
import android.widget.Button;
import android.widget.TextView;
import android.widget.Toast;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

public class MainActivity extends AppCompatActivity {

    private final int MY_PERMISSIONS_RECORD_AUDIO = 1;
    private static final String LOG_TAG = "WEKWS";
    private static final int SAMPLE_RATE = 16000;  // The sampling rate
    private static final int MAX_QUEUE_SIZE = 2500;  // 100 seconds audio, 1 / 0.04 * 100
    private static final List<String> resource = Arrays.asList("kws.ort");
    private static final String HINTS = "say \"Hey Snips\" to wakeup";

    private boolean startRecord = false;
    private AudioRecord record = null;
    private int miniBufferSize = 0;  // 1280 bytes 648 byte 40ms, 0.04s
    private final BlockingQueue<short[]> bufferQueue = new ArrayBlockingQueue<>(MAX_QUEUE_SIZE);

    public static void assetsInit(Context context) throws IOException {
        AssetManager assetMgr = context.getAssets();
        // Unzip all files in resource from assets to context.
        // Note: Uninstall the APP will remove the resource files in the context.
        for (String file : assetMgr.list("")) {
            if (resource.contains(file)) {
                File dst = new File(context.getFilesDir(), file);
                if (!dst.exists() || dst.length() == 0) {
                    Log.i(LOG_TAG, "Unzipping " + file + " to " + dst.getAbsolutePath());
                    InputStream is = assetMgr.open(file);
                    OutputStream os = new FileOutputStream(dst);
                    byte[] buffer = new byte[4 * 1024];
                    int read;
                    while ((read = is.read(buffer)) != -1) {
                        os.write(buffer, 0, read);
                    }
                    os.flush();
                }
            }
        }
    }

    @Override
    public void onRequestPermissionsResult(int requestCode,
                                           String[] permissions, int[] grantResults) {
        super.onRequestPermissionsResult(requestCode, permissions, grantResults);
        if (requestCode == MY_PERMISSIONS_RECORD_AUDIO) {
            if (grantResults.length > 0 && grantResults[0] == PackageManager.PERMISSION_GRANTED) {
                Log.i(LOG_TAG, "record permission is granted");
                initRecorder();
            } else {
                Toast.makeText(this, "Permissions denied to record audio", Toast.LENGTH_LONG).show();
                Button button = findViewById(R.id.button);
                button.setEnabled(false);
            }
        }
    }

    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_main);
        requestAudioPermissions();
        try {
            assetsInit(this);
        } catch (IOException e) {
            Log.e(LOG_TAG, "Error process asset files to file path");
        }

        TextView textView = findViewById(R.id.textView);
        textView.setText(HINTS);
        Spot.init(getFilesDir().getPath());

        Button button = findViewById(R.id.button);
        button.setText("Start Record");
        button.setOnClickListener(view -> {
            if (!startRecord) {
                startRecord = true;
                startRecordThread();
                startAcceptWaveThread();
                startSpotThread();
                button.setText("Stop Record");
            } else {
                startRecord = false;
                Spot.setInputFinished();
                button.setText("Start Record");
            }
        });
    }

    private void requestAudioPermissions() {
        if (ContextCompat.checkSelfPermission(this, Manifest.permission.RECORD_AUDIO)
                != PackageManager.PERMISSION_GRANTED) {
            ActivityCompat.requestPermissions(this,
                    new String[]{Manifest.permission.RECORD_AUDIO},
                    MY_PERMISSIONS_RECORD_AUDIO);
        } else {
            initRecorder();
        }
    }

    private void initRecorder() {
        // buffer size in bytes 1280
        miniBufferSize = AudioRecord.getMinBufferSize(SAMPLE_RATE,
                AudioFormat.CHANNEL_IN_MONO,
                AudioFormat.ENCODING_PCM_16BIT);
        if (miniBufferSize == AudioRecord.ERROR || miniBufferSize == AudioRecord.ERROR_BAD_VALUE) {
            Log.e(LOG_TAG, "Audio buffer can't initialize!");
            return;
        }
        if (ActivityCompat.checkSelfPermission(this, Manifest.permission.RECORD_AUDIO) != PackageManager.PERMISSION_GRANTED) {
            // TODO: Consider calling
            //    ActivityCompat#requestPermissions
            // here to request the missing permissions, and then overriding
            //   public void onRequestPermissionsResult(int requestCode, String[] permissions,
            //                                          int[] grantResults)
            // to handle the case where the user grants the permission. See the documentation
            // for ActivityCompat#requestPermissions for more details.
            return;
        }
        record = new AudioRecord(MediaRecorder.AudioSource.DEFAULT,
                SAMPLE_RATE,
                AudioFormat.CHANNEL_IN_MONO,
                AudioFormat.ENCODING_PCM_16BIT,
                miniBufferSize);
        if (record.getState() != AudioRecord.STATE_INITIALIZED) {
            Log.e(LOG_TAG, "Audio Record can't initialize!");
            return;
        }
        Log.i(LOG_TAG, "Record init okay");
    }

    private void startRecordThread() {
        new Thread(() -> {
            VoiceRectView voiceView = findViewById(R.id.voiceRectView);
            record.startRecording();
            Process.setThreadPriority(Process.THREAD_PRIORITY_AUDIO);
            while (startRecord) {
                short[] buffer = new short[miniBufferSize / 2];
                int read = record.read(buffer, 0, buffer.length);
                voiceView.add(calculateDb(buffer));
                try {
                    if (AudioRecord.ERROR_INVALID_OPERATION != read) {
                        bufferQueue.put(buffer);
                    }
                } catch (InterruptedException e) {
                    Log.e(LOG_TAG, e.getMessage());
                }
                Button button = findViewById(R.id.button);
                if (!button.isEnabled() && startRecord) {
                    runOnUiThread(() -> button.setEnabled(true));
                }
            }
            record.stop();
            voiceView.zero();
        }).start();
    }

    private double calculateDb(short[] buffer) {
        double energy = 0.0;
        for (short value : buffer) {
            energy += value * value;
        }
        energy /= buffer.length;
        energy = (10 * Math.log10(1 + energy)) / 100;
        energy = Math.min(energy, 1.0);
        return energy;
    }

    private void startAcceptWaveThread() {
        new Thread(() -> {
            // Send all data
            while (startRecord || bufferQueue.size() > 0) {
                try {
                    short[] data = bufferQueue.take();
                    // 1. add data to C++ interface
                    Spot.acceptWaveform(data);
                    // 2. get partial result
                    runOnUiThread(() -> {
                        TextView textView = findViewById(R.id.textView);
                        textView.setText(Spot.getResult());
                    });
                } catch (InterruptedException e) {
                    Log.e(LOG_TAG, e.getMessage());
                }
            }
        }).start();
    }

    private void startSpotThread() {
        new Thread(() -> {
            Spot.reset();
            // Send all data
            while (startRecord) {
                Spot.startSpot();
                Log.i(LOG_TAG, Spot.getResult());
            }
        }).start();
    }
}
